import numpy as np
import cv2
import glob
import math
import pandas as pd
import os
import sys
sys.path.append('././')
from kinematic import *
import transforms3d as tfs
# 标定板 9*6 25mm
chess_board_x_num = 10 - 1
chess_board_y_num = 7 - 1
chess_board_len = 25

# 相机内参矩阵和畸变系数
cam_mtx = np.array( [[670.92263044  ,      0.        ,   358.82135372],
                    [  0.           ,    673.8728403 ,   224.19994634],
                    [  0.           ,      0.        ,     1.        ]], dtype=np.float64)

cam_dist = np.array([-0.19770613, -0.19343703 ,-0.00058339, -0.00063087 , 0.07889203], dtype=np.float64)



count = 0

def get_RT_from_chessboard(img_path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist):
    '''
    #用来从棋盘格图片得到相机外参
    :param img_path: 读取图片路径
    :param chess_board_x_num: 棋盘格x方向格子数
    :param chess_board_y_num: 棋盘格y方向格子数
    :param chess_board_len: 单位棋盘格长度,mm
    :param cam_mtx: 相机内参
    :param cam_dist: 相机畸变参数
    :return: 相机外参
    '''
    global count
    
    # termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    # 标定板世界坐标
    objp = np.zeros((chess_board_y_num*chess_board_x_num,3), np.float32)
    for m in range(chess_board_y_num):
        for n in range(chess_board_x_num):
            objp[m*chess_board_x_num + n] = [n*chess_board_len, m*chess_board_len, 0]
    # print('obj:\n',objp)
    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.

    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Find the chess board corners
    ret, corners = cv2.findChessboardCorners(gray, (chess_board_x_num, chess_board_y_num), None)

    # If found, add object points, image points (after refining them)
    if ret == True:
        objpoints.append(objp)
        corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        draw_p=cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners, ret)
        cv2.imwrite('calibration/handeye_clib/drawChessboardCorners/'+str(count)+'.jpg', draw_p)
        count+=1
        # Draw and display the corners
        # cv2.namedWindow('img', 0)
        # cv2.resizeWindow('img', int(2560*0.6), int(1440*0.6))
        # cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners2, ret)
        # cv2.imshow('img', img)
        # cv2.waitKey(10)
        # input("请输入任意字符继续：")
        retval, rvec, tvec = cv2.solvePnP(objpoints[0], imgpoints[0], cam_mtx, cam_dist)
        # print(rvec.reshape((1,3)))
        # 反投影误差
        total_error = 0
        imgpoints2, _ = cv2.projectPoints(objpoints[0], rvec, tvec, cam_mtx, cam_dist)
        error = cv2.norm(imgpoints[0], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        total_error += error
        print("SolvePnP error: ", total_error / len(objpoints))

        RT = np.column_stack(((cv2.Rodrigues(rvec))[0],tvec))
        RT = np.row_stack((RT, np.array([0, 0, 0, 1])))
        


        return True,RT
    return False,None


def get_mtx(path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist):
    img_path =path+ "/*.jpg"
    coord_file_path = path + '/机械臂末端位姿.xlsx' #从记录文件读取机器人六个位姿
    df = pd.read_excel(coord_file_path)
    RT_all_end_to_base_1=[]
    RT_all_chess_to_cam_1=[]
    images = glob.glob(img_path)
    for row,fname in enumerate(images):
        print('第'+str(row)+'张图片:')
        ret,RT_chess_to_cam = get_RT_from_chessboard(fname, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist)
        if ret:
            RT_all_chess_to_cam_1.append(RT_chess_to_cam)
            print('solvePnP成功:')
            print(RT_chess_to_cam)

            joint_angles=[df.iloc[row, 0], df.iloc[row, 1], df.iloc[row, 2], df.iloc[row, 3], df.iloc[row, 4]]
            RT_end_to_base = forward_kinematic(np.deg2rad(joint_angles))
            RT_all_end_to_base_1.append(RT_end_to_base)
            print('正解成功:')
            print(RT_end_to_base)
            print('')
    
    return RT_all_end_to_base_1,RT_all_chess_to_cam_1

def get_matrix_eular_radu(x,y,z,rx,ry,rz):
    rmat = tfs.euler.euler2mat(math.radians(rx),math.radians(ry),math.radians(rz))
    rmat = tfs.affines.compose(np.squeeze(np.asarray((x,y,z))), rmat, [1, 1, 1])
    return rmat


def skew(v):
    return np.array([[0,-v[2],v[1]],
                     [v[2],0,-v[0]],
                     [-v[1],v[0],0]])


def rot2quat_minimal(m):
    quat =  tfs.quaternions.mat2quat(m[0:3,0:3])
    return quat[1:]


def quatMinimal2rot(q):
    p = np.dot(q.T,q)
    w = np.sqrt(np.subtract(1,p[0][0]))
    return tfs.quaternions.quat2mat([w,q[0],q[1],q[2]])



def hangeye_calibration():
    print("开始手眼标定\n")
    path = "calibration/handeye_clib/hand_eye"

    RT_all_end_to_base,RT_all_chess_to_cam=get_mtx(path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist)

    RT_chess_to_base_1=np.array([[0,-1,0,500],
                               [-1,0,0,100],
                               [0,0,-1,-10],
                               [0,0,0,1]
                               ])
    RT_chess_to_base_2=np.array([[0,-1,0,500],
                               [-1,0,0,100],
                               [0,0,-1,-10],
                               [0,0,0,1]
                               ])
    RT_all_cam_to_end=[]
    print(RT_all_end_to_base)
    #第一个位姿
    for i in range(4):
        RT_all_end_to_base_inv=np.linalg.inv(RT_all_end_to_base[i])
        RT_all_chess_to_cam_inv=np.linalg.inv(RT_all_chess_to_cam[i])

        RT_cam_to_end=RT_all_end_to_base_inv@RT_chess_to_base_1@RT_all_chess_to_cam_inv
        RT_all_cam_to_end.append(RT_cam_to_end)
    #第二个位姿
    
    for i in range(4,8):
        RT_all_end_to_base_inv=np.linalg.inv(RT_all_end_to_base[i])
        RT_all_chess_to_cam_inv=np.linalg.inv(RT_all_chess_to_cam[i])

        RT_cam_to_end=RT_all_end_to_base_inv@RT_chess_to_base_2@RT_all_chess_to_cam_inv
        RT_all_cam_to_end.append(RT_cam_to_end)
    RT_all_cam_to_end=np.array(RT_all_cam_to_end)
    

    # 初始化一个全零矩阵，用于存储矩阵的累加结果
    sum_matrix = np.zeros((4, 4), dtype=np.float64)

    # 累加所有矩阵
    for matrix in RT_all_cam_to_end[4:8,:]:
        sum_matrix += matrix

    # 计算平均矩阵
    average_matrix = sum_matrix / 4
    print('RT_all_cam_to_end:\n',RT_all_cam_to_end)
    print("RT_cam_to_end（"+str(len(RT_all_cam_to_end))+"次平均后）:\n",average_matrix)
    np.save('calibration/handeye_clib/eye_to_end.npy', average_matrix)
    

hangeye_calibration()


